﻿#pragma once

#include <fstream>
#include <iostream>
#include <learn/check.h>
#include <string>
#include <vector>

#pragma pack(push, 1)
struct TGA_Header
{
    std::uint8_t idlength{};
    std::uint8_t colormaptype{};
    std::uint8_t datatypecode{};
    std::uint16_t colormaporigin{};
    std::uint16_t colormaplength{};
    std::uint8_t colormapdepth{};
    std::uint16_t x_origin{};
    std::uint16_t y_origin{};
    std::uint16_t width{};
    std::uint16_t height{};
    std::uint8_t bitsperpixel{};
    std::uint8_t imagedescriptor{};
};
#pragma pack(pop)

struct TGA_Color
{
    std::uint8_t bgra[4];
    std::uint8_t bytespp; // 深度

    TGA_Color(std::uint8_t R = 0, std::uint8_t G = 0, std::uint8_t B = 0, std::uint8_t A = 255)
        : bgra{B, G, R, A}, bytespp(4)
    {
    }

    TGA_Color(const std::uint8_t *p, std::uint8_t bpp) : bytespp(bpp)
    {
        for (std::int32_t i = bpp; i--; bgra[i] = p[i])
            ;
    }

    TGA_Color(const TGA_Color &other)
        : bgra{other.bgra[0], other.bgra[1], other.bgra[2], other.bgra[3]}, bytespp(other.bytespp)
    {
    }

    TGA_Color &operator=(const TGA_Color &other)
    {
        bgra[0] = other.bgra[0];
        bgra[1] = other.bgra[1];
        bgra[2] = other.bgra[2];
        bgra[3] = other.bgra[3];
        bytespp = other.bytespp;
        return *this;
    }

    bool operator==(const TGA_Color &other) const
    {
        for (std::int32_t i = bytespp; i--;)
            if (bgra[i] != other.bgra[i])
                return false;
        return true;
    }

    bool operator!=(const TGA_Color &other) const
    {
        return !(*this == other);
    }
};

struct TGA_Image
{
    enum Format
    {
        GrayScale = 1,
        RGB = 3,
        RGBA = 4
    };

  private:
    struct _Impl
    {
        std::vector<std::uint8_t> data;

        int width{};
        int height{};
        Format bpp{RGBA};

        _Impl() = default;
        _Impl(int width, int height, Format bpp) : data(width * height * bpp, 0), width(width), height(height), bpp(bpp)
        {
        }
        _Impl(_Impl &&) = delete;
    };

    std::shared_ptr<_Impl> m_impl;

  public:
    TGA_Image() : m_impl(std::make_shared<_Impl>())
    {
    }

    TGA_Image(int width, int height, Format bpp) : m_impl(std::make_shared<_Impl>(width, height, bpp))
    {
    }

    _Impl *operator->() const noexcept
    {
        return m_impl.get();
    }

    bool read_tga_file(std::string filename)
    {
        std::ifstream in;
        in.open(filename, std::ios::binary);
        if (!in.is_open())
        {
            std::cerr << "can't open file " << filename << "\n";
            in.close();
            return false;
        }
        TGA_Header header;
        in.read(reinterpret_cast<char *>(&header), sizeof(header));
        if (!in.good())
        {
            in.close();
            std::cerr << "an error occured while reading the header\n";
            return false;
        }
        m_impl->width = header.width;
        m_impl->height = header.height;
        m_impl->bpp = Format(header.bitsperpixel >> 3);
        if (m_impl->width <= 0 || m_impl->height <= 0 ||
            (m_impl->bpp != GrayScale && m_impl->bpp != RGB && m_impl->bpp != RGBA))
        {
            in.close();
            std::cerr << "bad bpp (or width/height) value\n";
            return false;
        }
        std::size_t nbytes = m_impl->bpp * m_impl->width * m_impl->height;
        m_impl->data = std::vector<std::uint8_t>(nbytes, 0);
        if (3 == header.datatypecode || 2 == header.datatypecode)
        {
            in.read(reinterpret_cast<char *>(m_impl->data.data()), nbytes);
            if (!in.good())
            {
                in.close();
                std::cerr << "an error occured while reading the data\n";
                return false;
            }
        }
        else if (10 == header.datatypecode || 11 == header.datatypecode)
        {
            if (!load_rle_data(in))
            {
                in.close();
                std::cerr << "an error occured while reading the data\n";
                return false;
            }
        }
        else
        {
            in.close();
            std::cerr << "unknown file format " << (int)header.datatypecode << "\n";
            return false;
        }
        if (header.imagedescriptor & 0x20)
            flip_vertically();
        if (header.imagedescriptor & 0x10)
            flip_horizontally();
        std::cerr << m_impl->width << "x" << m_impl->height << "/" << m_impl->bpp * 8 << "\n";
        in.close();
        return true;
    }

    bool write_tga_file(const std::string filename, const bool vflip = true, const bool rle = true) const
    {
        constexpr std::uint8_t developer_area_ref[4] = {0, 0, 0, 0};
        constexpr std::uint8_t extension_area_ref[4] = {0, 0, 0, 0};
        constexpr std::uint8_t footer[18] = {'T', 'R', 'U', 'E', 'V', 'I', 'S', 'I', 'O',
                                             'N', '-', 'X', 'F', 'I', 'L', 'E', '.', '\0'};
        std::ofstream out;
        out.open(filename, std::ios::binary);
        if (!out.is_open())
        {
            std::cerr << "can't open file " << filename << "\n";
            out.close();
            return false;
        }
        TGA_Header header;
        header.bitsperpixel = m_impl->bpp << 3;
        header.width = m_impl->width;
        header.height = m_impl->height;
        header.datatypecode = (m_impl->bpp == GrayScale ? (rle ? 11 : 3) : (rle ? 10 : 2));
        header.imagedescriptor = vflip ? 0x00 : 0x20; // top-left or bottom-left origin
        out.write(reinterpret_cast<const char *>(&header), sizeof(header));
        if (!out.good())
        {
            out.close();
            std::cerr << "can't dump the tga file\n";
            return false;
        }
        if (!rle)
        {
            out.write(reinterpret_cast<const char *>(m_impl->data.data()),
                      m_impl->width * m_impl->height * m_impl->bpp);
            if (!out.good())
            {
                std::cerr << "can't unload raw data\n";
                out.close();
                return false;
            }
        }
        else if (!unload_rle_data(out))
        {
            out.close();
            std::cerr << "can't unload rle data\n";
            return false;
        }
        out.write(reinterpret_cast<const char *>(developer_area_ref), sizeof(developer_area_ref));
        if (!out.good())
        {
            std::cerr << "can't dump the tga file\n";
            out.close();
            return false;
        }
        out.write(reinterpret_cast<const char *>(extension_area_ref), sizeof(extension_area_ref));
        if (!out.good())
        {
            std::cerr << "can't dump the tga file\n";
            out.close();
            return false;
        }
        out.write(reinterpret_cast<const char *>(footer), sizeof(footer));
        if (!out.good())
        {
            std::cerr << "can't dump the tga file\n";
            out.close();
            return false;
        }
        out.close();
        return true;
    }

    void flip_horizontally()
    {
        int half = m_impl->width >> 1;
        for (int i = 0; i < half; i++)
            for (int j = 0; j < m_impl->height; j++)
                for (int b = 0; b < m_impl->bpp; b++)
                    std::swap(m_impl->data[(i + j * m_impl->width) * m_impl->bpp + b],
                              m_impl->data[(m_impl->width - 1 - i + j * m_impl->width) * m_impl->bpp + b]);
    }

    void flip_vertically()
    {
        int half = m_impl->height >> 1;
        for (int i = 0; i < m_impl->width; i++)
            for (int j = 0; j < half; j++)
                for (int b = 0; b < m_impl->bpp; b++)
                    std::swap(m_impl->data[(i + j * m_impl->width) * m_impl->bpp + b],
                              m_impl->data[(i + (m_impl->height - 1 - j) * m_impl->width) * m_impl->bpp + b]);
    }

    TGA_Color get(int x, int y) const
    {
        ASSERT_INFO(x >= 0 && x < m_impl->width && y >= 0 && y < m_impl->height, "pixel out of range");
        return TGA_Color(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, m_impl->bpp);
    }

    void set(int x, int y, const TGA_Color &c)
    {
        ASSERT_INFO(x >= 0 && x < m_impl->width && y >= 0 && y < m_impl->height, "pixel out of range");
        std::memcpy(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, c.bgra, m_impl->bpp);
    }

    void resize(int width, int height)
    {
        m_impl->data.resize(width * height * m_impl->bpp);
        m_impl->width = width;
        m_impl->height = height;
    }

  private:
    bool load_rle_data(std::ifstream &in)
    {
        std::size_t pixelcount = m_impl->width * m_impl->height;
        std::size_t currentpixel = 0;
        std::size_t currentbyte = 0;
        TGA_Color colorbuffer;
        do
        {
            std::uint8_t chunkheader = 0;
            chunkheader = in.get();
            if (!in.good())
            {
                std::cerr << "an error occured while reading the data\n";
                return false;
            }
            if (chunkheader < 128)
            {
                chunkheader++;
                for (int i = 0; i < chunkheader; i++)
                {
                    in.read(reinterpret_cast<char *>(colorbuffer.bgra), m_impl->bpp);
                    if (!in.good())
                    {
                        std::cerr << "an error occured while reading the header\n";
                        return false;
                    }
                    for (int t = 0; t < m_impl->bpp; t++)
                        m_impl->data[currentbyte++] = colorbuffer.bgra[t];
                    currentpixel++;
                    if (currentpixel > pixelcount)
                    {
                        std::cerr << "Too many pixels read\n";
                        return false;
                    }
                }
            }
            else
            {
                chunkheader -= 127;
                in.read(reinterpret_cast<char *>(colorbuffer.bgra), m_impl->bpp);
                if (!in.good())
                {
                    std::cerr << "an error occured while reading the header\n";
                    return false;
                }
                for (int i = 0; i < chunkheader; i++)
                {
                    for (int t = 0; t < m_impl->bpp; t++)
                        m_impl->data[currentbyte++] = colorbuffer.bgra[t];
                    currentpixel++;
                    if (currentpixel > pixelcount)
                    {
                        std::cerr << "Too many pixels read\n";
                        return false;
                    }
                }
            }
        } while (currentpixel < pixelcount);
        return true;
    }

    bool unload_rle_data(std::ofstream &out) const
    {
        const std::uint8_t max_chunk_length = 128;
        std::size_t npixels = m_impl->width * m_impl->height;
        std::size_t curpix = 0;
        while (curpix < npixels)
        {
            std::size_t chunkstart = curpix * m_impl->bpp;
            std::size_t curbyte = curpix * m_impl->bpp;
            std::uint8_t run_length = 1;
            bool raw = true;
            while (curpix + run_length < npixels && run_length < max_chunk_length)
            {
                bool succ_eq = true;
                for (int t = 0; succ_eq && t < m_impl->bpp; t++)
                    succ_eq = (m_impl->data[curbyte + t] == m_impl->data[curbyte + t + m_impl->bpp]);
                curbyte += m_impl->bpp;
                if (1 == run_length)
                    raw = !succ_eq;
                if (raw && succ_eq)
                {
                    run_length--;
                    break;
                }
                if (!raw && !succ_eq)
                    break;
                run_length++;
            }
            curpix += run_length;
            out.put(raw ? run_length - 1 : run_length + 127);
            if (!out.good())
            {
                std::cerr << "can't dump the tga file\n";
                return false;
            }
            out.write(reinterpret_cast<const char *>(m_impl->data.data() + chunkstart),
                      (raw ? run_length * m_impl->bpp : m_impl->bpp));
            if (!out.good())
            {
                std::cerr << "can't dump the tga file\n";
                return false;
            }
        }
        return true;
    }
};